package cn.edu.hitsz.compiler.parser;

import cn.edu.hitsz.compiler.ir.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;

/** 中间代码优化器 */
public class IROptimizer {
    private List<Instruction> inst = new ArrayList<Instruction>();

    /** 基本块入口处的活跃变量集合 */
    private HashSet<IRVariable> in = new HashSet<IRVariable>();

    /** 基本块出口处的活跃变量集合 */
    private HashSet<IRVariable> out = new HashSet<IRVariable>();

    /** 在基本块中被定义，且在定义前在块中未被引用的变量集合 */
    private HashSet<IRVariable> def = new HashSet<IRVariable>();

    /** 在基本块中被引用，且在引用前未在块中被定义过的(变量名,变量值)键值对集合  */
    private HashMap<IRVariable,Integer> use = new HashMap<IRVariable, Integer>();


    public IROptimizer(List<Instruction> inst) {
        this.inst = inst;
    }

    /** 活跃变量分析 */
    public void analyzer() {
        /** 目前的编译器结构比较简单，只有一个基本块，因此in一定为空 */
        /** 返回指令只能且一定出现在最后一条 否则过不了语法分析 */

        // 找到返回指令（最后一条指令） 将返回变量加入out
        Instruction retIR = inst.get(inst.size()-1);

        if(retIR.getKind() == InstructionKind.RET){
            out.add((IRVariable)retIR.getReturnValue());
        }
    }

    /** 将IRValue与use表中的变量映射获取值
     *  如果是IRImmediate 则直接返回值
     *  如果是IRVariable 则根据use表里的值返回
     *  如果表中没有，则说明代码中使用了未定义值的变量运算 这是不合法的
     * */
    public int mapping(IRValue ir){
        int val;
        if(ir.isIRVariable()){
            if(use.containsKey(ir)){
                val = use.get(ir);
            }
            else {
                throw new RuntimeException("IROptimizer error!");
            }
        }
        else {
            val = Integer.parseInt(ir.toString());
        }
        return val;
    }

    /** 从前往后数据流分析. 对单基本块的常量进行复写传播优化、无用指令删除 */
    public void run() {
        // 先进行活跃变量分析
        this.analyzer();
        int idx = 0;
        boolean flag = false;

        while (true){
            Instruction ir = inst.get(idx);
            switch (ir.getKind()){
                case MOV -> {
                    IRVariable res = ir.getResult();
                    IRValue from = ir.getFrom();

                    // 删除无用指令并进行常量传播或复写
                    inst.remove(idx);
                    use.put(res, mapping(from));
                }

                case RET -> {
                    // 删除无用指令
                    inst.remove(idx);

                    IRValue returnVal = ir.getReturnValue();

                    // 尝试进行常量传播
                    inst.add(Instruction.createRet(IRImmediate.of(mapping(returnVal))));

                    // 标志优化结束
                    flag = true;
                }

                case ADD -> {
                    IRVariable result = ir.getResult();
                    IRValue lhs = ir.getLHS();
                    IRValue rhs = ir.getRHS();

                    // 删除无用指令并进行常量传播、折叠
                    inst.remove(idx);
                    use.put(result, mapping(lhs) + mapping(rhs));
                }

                case SUB -> {
                    IRVariable result = ir.getResult();
                    IRValue lhs = ir.getLHS();
                    IRValue rhs = ir.getRHS();

                    // 删除无用指令并进行常量传播、折叠
                    inst.remove(idx);
                    use.put(result, mapping(lhs) - mapping(rhs));
                }

                case MUL -> {
                    IRVariable result = ir.getResult();
                    IRValue lhs = ir.getLHS();
                    IRValue rhs = ir.getRHS();

                    // 删除无用指令并进行常量传播、折叠
                    inst.remove(idx);
                    use.put(result, mapping(lhs) * mapping(rhs));
                }
            }
            // 数据流分析结束
            if(flag){
                break;
            }
        }
        System.out.println(use);
    }

    /** 获得IR优化结果 */
    public List<Instruction> getIR() {
        return inst;
    }
}
